#!/usr/bin/env python
# -*- encoding: utf-8 -*-
"""
@File    :   PSO.py    
@Contact :   ruaqy@qq.com
@License :   (C)Copyright 2022-~, GPL 2.0

@Modify Time      @Author    @Version    @Description
------------      -------    --------    -----------
2022/7/2 15:17   rqy        1.0         None
"""
import numpy as np
import matplotlib.pyplot as plt


class Particle:
    def __init__(self, v, x, w, c1, c2):
        self.v = v
        self.x = x
        self.p_p_best = x
        self.w = w
        self.c1 = c1
        self.c2 = c2

        self.best_adaption = None

    def v_update(self, p_g_best):
        r1 = np.random.random()
        r2 = np.random.random()
        self.v = self.w * self.v + self.c1 * r1 * (self.p_p_best - self.x) + self.c2 * r2 * (p_g_best - self.x)

    def x_update(self):
        self.x = self.x + self.v

    def get_adaption(self, adaption):
        adaption_value = adaption(self.x)
        if self.best_adaption:
            if self.best_adaption < adaption_value:
                self.p_p_best = self.x
                self.best_adaption = adaption_value
        else:
            self.best_adaption = adaption_value
        return adaption_value


# adaption function
def adaption_function(x):
    # return -np.sum(np.square(x-np.array([0.5, 0.5])))
    return -np.sum(np.square(np.power(x, 3)-np.array([0.5, 0.5])))


def main():
    N = 100
    T = 100
    D = 2
    w = 0.8
    c1 = 1
    c2 = 1

    #  initial particles value
    particles = [Particle(np.random.randn(D), np.random.randn(D), w, c1, c2) for i in range(N)]
    adaption = [i.get_adaption(adaption_function) for i in particles]
    p_g_best = particles[np.argmax(adaption)].x
    g_best_adaption_value = adaption_function(p_g_best)

    # usage for figure
    best_adaption_list = []

    for k in range(T):
        # update the x,v of particles and calculate the adaption
        [i.v_update(p_g_best) for i in particles]
        [i.x_update() for i in particles]
        adaption = [i.get_adaption(adaption_function) for i in particles]

        # calculate the best adaption of k times
        k_g_best = particles[np.argmax(adaption)].x
        k_best_adaption_value = adaption_function(k_g_best)

        if k_best_adaption_value > g_best_adaption_value:
            g_best_adaption_value = k_best_adaption_value
            p_g_best = k_g_best

        # usage for figure
        plt.clf()
        plt.scatter([i.x[0] for i in particles], [i.x[1] for i in particles])
        plt.title('iter:{} best adaption:{}'.format(k, max(adaption)))
        plt.xlim(-2, 2)
        plt.ylim(-2, 2)
        plt.pause(0.1)

    plt.show()


if __name__ == '__main__':
    main()

